#!/usr/bin/env python

import asyncio
from mcp.server.fastmcp import FastMCP, Context
from mcp.client.session import Session
from mcp.client.stdlib.http import HttpMCPTransport  # 支持 Streamable HTTP 远程 MCP
from mcp.client.stdlib.sse import SSEClientTransport  # 如远程是 SSE server

mcp = FastMCP("LocalProxy")

async def get_remote_session():
    """
    select transport based on remote MCP server：
    - Streamable HTTP: HttpMCPTransport
    - SSE: SSEClientTransport
    """

    transport = HttpMCPTransport("http://localhost:4000/mcp")

    session = Session(transport)
    await session.initialize()
    return session


@mcp.tool()
async def remote_add(a: int, b: int, ctx: Context) -> str:
    """
    调用远程 MCP server 的 add 工具
    """
    session = await get_remote_session()

    # 调用远程工具
    remote_result = await session.call_tool(
        name="add",
        arguments={"a": a, "b": b}
    )

    await session.close()

    # remote_result.content 是 MCP 标准格式
    text = remote_result.content[0].text
    return f"Remote result: {text}"

@mcp.tool()
async def remote_multiply(a: int, b: int, ctx: Context) -> str:
    session = await get_remote_session()

    remote_result = await session.call_tool(
        name="multiply",
        arguments={"a": a, "b": b}
    )

    await session.close()

    text = remote_result.content[0].text
    return f"Remote result: {text}"


if __name__ == "__main__":
    mcp.run()
